package org.geekbang.time.mathbasicthought.lesson10_depth_first_traversal;

/**
 * 给你一个字典，请尝试实现其前缀树，包括树的构建和查询两个过程。
 * 这里，字典可以用字符串数组表示，每个字符串代表一个单词。
 */
public class PrefixTreeTest {

    private static String[] str = {"key", "value", "wuyanzhao", "abandon", "k"};
    private static TrieNode root = new TrieNode();

    public static void main(String[] args) {
        for (String key : str) {
            insert(key);
        }

        System.out.println("their result:" + search("their"));
        System.out.println("k result:" + search("k"));
        System.out.println("wuyanzhao result:" + search("wuyanzhao"));
        System.out.println("wy result:" + search("wy"));
    }

    public static boolean search(String key) {
        if (key == null || key.trim().length() == 0) {
            return false;
        }

        TrieNode searchNode = root;
        for (char ch : key.toCharArray()) {
            int index = ch - 'a';

            if (searchNode.children[index] == null)
                return false;

            searchNode = searchNode.children[index];
        }

        return (searchNode != null && searchNode.isWordEnd);
    }

    public static void insert(String key) {
        if (key == null || key.trim().length() == 0) {
            return;
        }

        TrieNode newNode = root;

        for (char ch : key.toCharArray()) {
            int index = ch - 'a';

            if (newNode.children[index] == null)
                newNode.children[index] = new TrieNode();

            newNode = newNode.children[index];
        }

        newNode.isWordEnd = true;
    }

    static class TrieNode {
        private boolean isWordEnd;

        private TrieNode[] children;

        public TrieNode() {
            isWordEnd = false;

            children = new TrieNode[26];
            for (TrieNode child : children) {
                child = null;
            }
        }
    }
}
